from importlib import import_module

from rocket.estimator.evaluator.data_evaluator import ClsDataEvaluator
from rocket.estimator.evaluator.grad_evaluator import SimpleGradEvaluator
from rocket.estimator.evaluator.loss_evaluator import ClsLossEvaluator
from rocket.estimator.hook.checkpoint import CheckpointHook
from rocket.estimator.hook.dispatcher import HookDispatcher
from rocket.estimator.hook.logger import LoggerHook
from rocket.estimator.hook.loss import LossPlotHook
from rocket.estimator.hook.metric import MetricHook
from rocket.estimator.hook.timer import IterTimerHook, EpochTimerHook


class Registry:

    def __init__(self):
        self.catalog = {}

    def register(self, cls):
        cls_name = cls.__name__
        self.catalog[cls_name] = cls

    def get(self, cls_name):
        return self.catalog[cls_name]

    def contains(self, cls_name):
        return cls_name in self.catalog


class ExternalRegistry(Registry):

    def get(self, cls_name):
        if self.contains(cls_name):
            return self.catalog[cls_name]
        else:
            return self._get_cls(cls_name)

    @staticmethod
    def _get_cls(cls_name):
        path = cls_name.split(".")
        module_name = ".".join(path[:-1])
        cls_name = path[-1]

        module = import_module(module_name)
        cls = getattr(module, cls_name)

        return cls


REGISTRY = ExternalRegistry()
REGISTRY.register(ClsDataEvaluator)
REGISTRY.register(ClsLossEvaluator)
REGISTRY.register(SimpleGradEvaluator)
REGISTRY.register(HookDispatcher)
REGISTRY.register(LoggerHook)
REGISTRY.register(MetricHook)
REGISTRY.register(IterTimerHook)
REGISTRY.register(EpochTimerHook)
REGISTRY.register(CheckpointHook)
REGISTRY.register(LossPlotHook)
